/*
 * Copyright (c) 2017-present, CV4J Contributors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.cv4j.core.binary;

import com.cv4j.core.datamodel.ByteProcessor;
import com.cv4j.core.datamodel.IntIntegralImage;

/**
 * Threshold
 *
 * @since 2021-07-22
 */
public class Threshold {
    /**
     * binary image
     */
    public static final int METHOD_THRESH_BINARY = 0;
    /**
     * invert binary image
     */
    public static final int METHOD_THRESH_BINARY_INV = 1;

    /**
     * it is not reasonable method to convert binary image
     */
    public static final int THRESH_MEANS = 1;
    /**
     * it is popular binary method in OPENCV and MATLAB
     */
    public static final int THRESH_OTSU = 2;
    /**
     * histogram statistic threshold method
     */
    public static final int THRESH_TRIANGLE = 3;
    /**
     * based on 1D mean shift, CV4J custom binary method, sometimes it is very slow...
     */
    public static final int THRESH_MEANSHIFT = 4;
    /**
     * based local mean threshold method, CV4J custom binary method, sometimes it is very slow...
     */
    public static final int ADAPTIVE_C_MEANS_THRESH = 5;
    /**
     * it is not reasonable method to convert binary image
     */
    public static final int THRESH_VALUE = -1;

    /**
     * process
     *
     * @param gray - gray image data, ByteProcessor type
     * @param type - binary segmentation method, int
     */
    public void process(ByteProcessor gray, int type) {
        process(gray, type, METHOD_THRESH_BINARY, 0);
    }

    /**
     * adaptiveThresh
     *
     * @param gray gray
     * @param type type
     * @param blockSize blockSize
     * @param constant constant
     * @param method method
     */
    public void adaptiveThresh(ByteProcessor gray, int type, int blockSize, int constant, int method) {
        int width = gray.getWidth();
        int height = gray.getHeight();

        // 图像灰度化

        // per-calculate integral image
        IntIntegralImage grayii = new IntIntegralImage();
        byte[] binaryData = gray.getGray();
        grayii.setImage(binaryData);
        grayii.calculate(width, height);
        int yr = blockSize;
        int xr = blockSize;
        int index = 0;
        int size = (yr * 2 + 1) * (xr * 2 + 1);
        for (int row = 0; row < height; row++) {
            for (int col = 0; col < width; col++) {
                index = row * width + col;

                // 计算均值
                int sr = grayii.getBlockSum(col, row, (yr * 2 + 1), (xr * 2 + 1));
                int mean = sr / size;
                int pixel = binaryData[index] & 0xff;

                // 二值化
                if (pixel > (mean - constant)) {
                    binaryData[row * width + col] = (byte) 255;
                } else {
                    binaryData[row * width + col] = (byte) 0;
                }
            }
        }
    }

    /**
     * process
     *
     * @param gray - gray image data, ByteProcessor type
     * @param type - binary segmentation method, int
     * @param method -method
     * @param thresh - threshold value you are going to use it if type = 0;
     */
    public void process(ByteProcessor gray, int type, int method, int thresh) {
        int tvalue = thresh;
        if (type == THRESH_MEANS) {
            tvalue = getMeanThreshold(gray);
        } else if (type == THRESH_OTSU) {
            tvalue = getOTSUThreshold(gray);
        } else if (type == THRESH_TRIANGLE) {
            tvalue = getTriangleThreshold(gray);
        } else if (type == THRESH_MEANSHIFT) {
            tvalue = shift(gray);
        } else if (type == THRESH_VALUE) {
            tvalue = thresh;
        } else {
            tvalue = thresh;
        }
        byte[] data = gray.getGray();
        int c1 = 0;
        for (int i = 0; i < data.length; i++) {
            c1 = data[i] & 0xff;
            if (c1<= tvalue) {
                data[i] = (method == METHOD_THRESH_BINARY_INV) ? (byte) 255 : (byte) 0;
            } else {
                data[i] = (method == METHOD_THRESH_BINARY_INV) ? (byte) 0 : (byte) 255;
            }
        }
    }

    private int getMeanThreshold(ByteProcessor gray) {
        byte[] data = gray.getGray();
        int sum = 0;
        for (int i = 0; i < data.length; i++) {
            sum += data[i] & 0xff;
        }
        return sum / data.length;
    }

    private int getOTSUThreshold(ByteProcessor gray) {
        // 获取直方图
        int[] histogram = new int[256];
        byte[] data = gray.getGray();
        int c1 = 0;
        for (int i = 0; i < data.length; i++) {
            c1 = data[i] & 0xff;
            histogram[c1]++;
        }
        // 图像二值化 - OTSU 阈值化方法
        double total = data.length;
        double[] variances = new double[256];
        for (int i = 0; i < variances.length; i++) {
            double bw = 0;
            double bmeans = 0;
            double bvariance = 0;
            double count = 0;
            for (int t = 0; t < i; t++) {
                count += histogram[t];
                bmeans += histogram[t] * t;
            }
            bw = count / total;
            bmeans = (count == 0) ? 0 : (bmeans / count);
            for (int t = 0; t < i; t++) {
                bvariance += (Math.pow((t - bmeans), 2) * histogram[t]);
            }
            bvariance = (count == 0) ? 0 : (bvariance / count);
            double fw = 0;
            double fmeans = 0;
            double fvariance = 0;
            count = 0;
            for (int t = i; t < histogram.length; t++) {
                count += histogram[t];
                fmeans += histogram[t] * t;
            }
            fw = count / total;
            fmeans = (count == 0) ? 0 : (fmeans / count);
            for (int t = i; t < histogram.length; t++) {
                fvariance += (Math.pow((t - fmeans), 2) * histogram[t]);
            }
            fvariance = (count == 0) ? 0 : (fvariance / count);
            variances[i] = bw * bvariance + fw * fvariance;
        }

        // find the minimum within class variance
        double min = variances[0];
        int threshold = 0;
        for (int m = 1; m < variances.length; m++) {
            if (min > variances[m]) {
                threshold = m;
                min = variances[m];
            }
        }
        return threshold;
    }

    private int getTriangleThreshold(ByteProcessor gray) {
        // 获取直方图
        int[] histogram = new int[256];
        byte[] data = gray.getGray();
        int c1 = 0;
        for (int i = 0; i < data.length; i++) {
            c1 = data[i] & 0xff;
            histogram[c1]++;
        }

        int leftBound = 0;
        int rightBound = 0;
        int maxInd = 0;
        int max = 0;
        int temp;
        boolean isflipped = false;
        int i1 = 0;
        int j1 = 0;
        int n1 = 256;

        // 找到最左边零的位置
        for (i1 = 0; i1 < n1; i1++) {
            if (histogram[i1] > 0) {
                leftBound = i1;
                break;
            }
        }
        // 位置再移动一个步长，即为最左侧零位置
        if (leftBound > 0) {
            leftBound--;
        }

        // 找到最右边零点位置
        for (i1 = n1 - 1; i1 > 0; i1--) {
            if (histogram[i1] > 0) {
                rightBound = i1;
                break;
            }
        }
        // 位置再移动一个步长，即为最右侧零位置
        if (rightBound < n1 - 1) {
            rightBound++;
        }

        // 在直方图上寻找最亮的点Hmax
        for (i1 = 0; i1 < n1; i1++) {
            if (histogram[i1] > max) {
                max = histogram[i1];
                maxInd = i1;
            }
        }

        // 如果最大值落在靠左侧这样就无法满足三角法求阈值，所以要检测是否最大值是否靠近左侧
        // 如果靠近左侧则通过翻转到右侧位置。
        if (maxInd - leftBound < rightBound - maxInd) {
            isflipped = true;
            i1 = 0;
            j1 = n1 - 1;
            while (i1 < j1) {
                // 左右交换
                temp = histogram[i1];
                histogram[i1] = histogram[j1];
                histogram[j1] = temp;
                i1++;
                j1--;
            }
            leftBound = n1 - 1 - rightBound;
            maxInd = n1 - 1 - maxInd;
        }

        // 计算求得阈值
        double thresh = leftBound;
        double a1;
        double b1;
        double dist = 0;
        double tempdist;
        a1 = max;
        b1 = leftBound - maxInd;
        for (i1 = leftBound + 1; i1 <= maxInd; i1++) {
            // 计算距离 - 不需要真正计算
            tempdist = a1 * i1 + b1 * histogram[i1];
            if (tempdist > dist) {
                dist = tempdist;
                thresh = i1;
            }
        }
        thresh--;

        // 对已经得到的阈值T,如果前面已经翻转了，则阈值要用255-T
        if (isflipped) {
            thresh = n1 - 1 - thresh;
        }

        return (int) thresh;
    }

    private int shift(ByteProcessor gray) {
        // find threshold
        int t1 = 127;
        int nt = 0;
        int m1 = 0;
        int m2 = 0;
        int sum1 = 0;
        int sum2 = 0;
        int count1 = 0;
        int count2 = 0;
        int[] data = gray.toInt(0);
        while (true) {
            for (int i = 0; i < data.length; i++) {
                if (data[i] > t1) {
                    sum1 += data[i];
                    count1++;
                } else {
                    sum2 += data[i];
                    count2++;
                }
            }
            m1 = sum1 / count1;
            m2 = sum2 / count2;

            sum1 = 0;
            sum2 = 0;
            count1 = 0;
            count2 = 0;
            nt = (m1 + m2) / 2;
            if (t1 == nt) {
                break;
            } else {
                t1 = nt;
            }
        }
        return t1;
    }
}
